##根据tick数据计算RSI

from typing import Dict, Optional


class RSICalculator:
    """
    同花顺风格 RSI 计算器（按股票代码维护状态）。

    计算公式：
        RSI(N) = 100 * SMA(UP, N, 1) / SMA(ABS(DIFF), N, 1)
        其中 SMA 的递推为：SMA_t = (1*X_t + (N-1)*SMA_{t-1}) / N；SMA_1 = X_1。

    使用方式：
        rsi_calc = RSICalculator(period=10)
        rsi = rsi_calc.calculate_rsi("513130.SH", last_price)

    特性：
    - 与 Wilder 平滑等价的递推系数，但采用同花顺/通达信的 SMA 初始条件（首值等于首个样本）。
    - 初始化更快：第二笔变动即可给出 RSI 值。
    - 每个 tick O(1) 更新，适合高频 tick 场景。
    """

    def __init__(self, period: int = 10):
        if period <= 0:
            raise ValueError("RSI period must be positive")
        self.period: int = period

        # 每个股票的状态
        self._last_price: Dict[str, float] = {}
        # 同花顺 RSI：分别维护 UP 与 ABS(DIFF) 的 SMA(N,1)
        self._sma_up: Dict[str, float] = {}
        self._sma_abs: Dict[str, float] = {}

    def reset(self, stock_code: Optional[str] = None) -> None:
        """重置指定股票或全部股票的 RSI 状态。"""
        if stock_code is None:
            self._last_price.clear()
            self._sma_up.clear()
            self._sma_abs.clear()
            return

        self._last_price.pop(stock_code, None)
        self._sma_up.pop(stock_code, None)
        self._sma_abs.pop(stock_code, None)

    def calculate_rsi(self, stock_code: str, last_price: Optional[float]) -> Optional[float]:
        """
        增量计算 RSI 并返回最新值；若初始化不足，返回 None。

        参数:
            stock_code: 股票代码
            last_price: 最新价格（None 时返回 None）
        返回:
            RSI 数值(0~100)，或 None（未完成初始化）
        """
        if last_price is None:
            return None

        # 第一次到价仅记录，不计算涨跌
        if stock_code not in self._last_price:
            self._last_price[stock_code] = float(last_price)
            return None

        previous_price = self._last_price[stock_code]
        change = float(last_price) - previous_price
        up = change if change > 0 else 0.0
        abs_change = abs(change)

        # 若尚未有 SMA 值，按同花顺定义：SMA_1 = X_1
        if stock_code not in self._sma_up or stock_code not in self._sma_abs:
            self._sma_up[stock_code] = up
            self._sma_abs[stock_code] = abs_change
            self._last_price[stock_code] = float(last_price)
            denom = self._sma_abs[stock_code]
            if denom == 0:
                return 100.0 if self._sma_up[stock_code] > 0 else 0.0
            return 100.0 * (self._sma_up[stock_code] / denom)

        # 同花顺 SMA(N,1) 增量更新
        sma_up = self._sma_up[stock_code]
        sma_abs = self._sma_abs[stock_code]
        n = self.period
        sma_up = (sma_up * (n - 1) + up) / n
        sma_abs = (sma_abs * (n - 1) + abs_change) / n
        self._sma_up[stock_code] = sma_up
        self._sma_abs[stock_code] = sma_abs
        self._last_price[stock_code] = float(last_price)

        if sma_abs == 0:
            return 100.0 if sma_up > 0 else 0.0

        return 100.0 * (sma_up / sma_abs)


class THSRSICalculator:
    """
    同花顺口径 RSI(N)：仅基于收盘价序列（CLOSE）计算：
        LC = 前一根K线收盘价
        UP_t   = max(CLOSE_t - LC, 0)
        ABS_t  = abs(CLOSE_t - LC)
        RSI(N) = 100 * SMA(UP, N, 1) / SMA(ABS, N, 1)

    说明：
    - 只能在“收盘价确认”时调用 update_with_close(...) 更新；tick 间不更新。
    - SMA 递推：SMA_t = (1*X_t + (N-1)*SMA_{t-1}) / N；首值取首个样本 X_1。
    - get_last_rsi 返回最近一次已确认K线的 RSI 值。
    """

    def __init__(self, period: int = 10):
        if period <= 0:
            raise ValueError("RSI period must be positive")
        self.period: int = period
        self._last_close: Dict[str, float] = {}
        self._sma_up: Dict[str, float] = {}
        self._sma_abs: Dict[str, float] = {}
        self._last_rsi: Dict[str, float] = {}

    def reset(self, stock_code: Optional[str] = None) -> None:
        if stock_code is None:
            self._last_close.clear()
            self._sma_up.clear()
            self._sma_abs.clear()
            self._last_rsi.clear()
            return
        self._last_close.pop(stock_code, None)
        self._sma_up.pop(stock_code, None)
        self._sma_abs.pop(stock_code, None)
        self._last_rsi.pop(stock_code, None)

    def update_with_close(self, stock_code: str, close_price: Optional[float]) -> Optional[float]:
        if close_price is None:
            return self._last_rsi.get(stock_code)
        close_price = float(close_price)

        if stock_code not in self._last_close:
            self._last_close[stock_code] = close_price
            return None

        change = close_price - self._last_close[stock_code]
        up = change if change > 0 else 0.0
        abs_change = abs(change)

        if stock_code not in self._sma_up or stock_code not in self._sma_abs:
            self._sma_up[stock_code] = up
            self._sma_abs[stock_code] = abs_change
        else:
            n = self.period
            self._sma_up[stock_code] = (self._sma_up[stock_code] * (n - 1) + up) / n
            self._sma_abs[stock_code] = (self._sma_abs[stock_code] * (n - 1) + abs_change) / n

        self._last_close[stock_code] = close_price

        denom = self._sma_abs[stock_code]
        if denom == 0:
            rsi = 100.0 if self._sma_up[stock_code] > 0 else 0.0
        else:
            rsi = 100.0 * (self._sma_up[stock_code] / denom)
        self._last_rsi[stock_code] = rsi
        return rsi

    def get_last_rsi(self, stock_code: str) -> Optional[float]:
        return self._last_rsi.get(stock_code)

    def get_preview_rsi(self, stock_code: str, last_price: Optional[float]) -> Optional[float]:
        """
        预览当前分钟的 RSI（不改变内部状态）。
        逻辑：以“上一已确认分钟收盘价”为 LC，使用当前最新价作为本分钟的临时 close，计算临时 UP/ABS，
        带入 SMA 递推公式得到临时 RSI。分钟结束时仍需调用 update_with_close 正式落库。
        """
        if last_price is None:
            return self._last_rsi.get(stock_code)
        if stock_code not in self._last_close:
            return None
        try:
            n = self.period
            up = max(float(last_price) - float(self._last_close[stock_code]), 0.0)
            ab = abs(float(last_price) - float(self._last_close[stock_code]))
            if stock_code not in self._sma_up or stock_code not in self._sma_abs:
                sma_up_prev = 0.0
                sma_ab_prev = 0.0
            else:
                sma_up_prev = float(self._sma_up[stock_code])
                sma_ab_prev = float(self._sma_abs[stock_code])
            sma_up_tmp = (sma_up_prev * (n - 1) + up) / n
            sma_ab_tmp = (sma_ab_prev * (n - 1) + ab) / n
            if sma_ab_tmp == 0:
                return 100.0 if sma_up_tmp > 0 else 0.0
            return 100.0 * (sma_up_tmp / sma_ab_tmp)
        except Exception:
            return self._last_rsi.get(stock_code)


__all__ = ["RSICalculator", "THSRSICalculator"]


